"""
用于混合式 ReAct + 计划-执行智能体的规划模块。
"""

import json
import re
from typing import Any, Dict, List, Optional
import time
from json_repair import repair_json
from langchain.schema import HumanMessage
from langgraph.store.memory import BaseStore
import uuid
from loguru import logger
import traceback

from model.memory_model import MemoryType
from model.planer_model import Plan, PlanType, PlanStep
from llm.llm_manager import get_llm_manager
from agent.agent_prompts import plan_prompt_template, plan_refinement_prompt_template


class Planner:
    """高级规划器，可为复杂查询创建执行计划。"""
    
    def __init__(self, memory_store: BaseStore):
        self.memory_store = memory_store
        self.memory_type = MemoryType.PLAN_MEMORY.value
        self.llm_manager = get_llm_manager()

    async def create_plan(
            self, query: str, available_tools: List[Dict],
            context: Optional[Dict[str, Any]] = None, session_id: str = None
    ) -> Plan:

        # 从记忆中获取类似的成功计划
        similar_plans = await self._get_similar_plans(query)
        
        # 创建计划提示词
        prompt = self._create_planning_prompt(query, available_tools, similar_plans, context)

        messages = [
            HumanMessage(content=prompt)
        ]
        response = await self.llm_manager.ainvoke(messages, session_id)
        plan_text = response.content

        # 解析计划
        plan = self._parse_plan(query, plan_text, available_tools)

        # 保存计划到memory
        await self._store_plan(plan)

        return plan
    
    async def refine_plan(
            self, plan: Plan, execution_results: List[Dict[str, Any]],
            current_context: Dict[str, Any], session_id: str = None
    ) -> Plan:
        """根据执行结果优化计划。"""

        prompt = self._create_refinement_prompt(plan, execution_results, current_context)

        messages = [
            HumanMessage(content=prompt)
        ]

        response = await self.llm_manager.ainvoke(messages, session_id)
        refinement_text = response.content
        
        # 解析优化内容并更新计划
        refined_plan = self._apply_refinements(plan, refinement_text)
        
        # 保存计划到memory
        await self._store_plan(refined_plan)

        return refined_plan
    
    async def _get_similar_plans(self, query: str, limit: int = 3) -> List[Plan]:
        """从记忆中获取类似的成功计划。"""
        similar_memories = self.memory_store.search(
            (self.memory_type, ),
            query=query,
            limit=limit
        )

        successful_plans = []
        for memory in similar_memories:
            plan = Plan(**json.loads(memory.value['plan']))
            successful_plans.append(plan)
        
        return successful_plans
    
    def _create_planning_prompt(self, query: str, available_tools: List[Dict],
                               similar_plans: List[Plan],
                               context: Optional[Dict[str, Any]]) -> str:
        """创建规划提示词."""
        similar_plans_text = ""
        if similar_plans:
            similar_plans_text = "\n\n过去类似的成功计划：\n"
            for i, plan in enumerate(similar_plans, 1):
                similar_plans_text += f"\nPlan {i}:\n"
                similar_plans_text += f"Query: {plan.query}\n"
                steps = [step.model_dump_json() for step in plan.steps]
                similar_plans_text += f"Steps: {json.dumps(steps, ensure_ascii=False, indent=4)}\n"
        
        context_text = ""
        if context:
            context_text = f"\n\n当前上下文：\n{json.dumps(context, ensure_ascii=False, indent=4)}"

        return plan_prompt_template.format(
            inputs={
                "query": query,
                "tools_description": available_tools,
                "similar_plans_text": similar_plans_text,
                "context_text": context_text,
            },
            remove_template_variables=True
        )

    def _create_refinement_prompt(self, plan: Plan, execution_results: List[Dict[str, Any]], 
                                 current_context: Dict[str, Any]) -> str:
        """创建用于完善计划的提示词"""
        
        results_summary = []
        for result in execution_results:
            step_id = result.get("step_id", "unknown")
            success = result.get("success", False)
            error = result.get("error", "")
            results_summary.append(f"Step {step_id}: {'SUCCESS' if success else 'FAILED'} - {error}")

        return plan_refinement_prompt_template.format(
            inputs={
                "goal": plan.goal,
                "plan_type": plan.plan_type.value,
                "steps": [step.__dict__ for step in plan.steps],
                "results_summary": '\n'.join(results_summary),
                "current_context": current_context
            },
            remove_template_variables=True
        )

    def _parse_plan(self, query: str, plan_text: str, available_tools: List[Dict]) -> Plan:
        """解析计划"""
        try:
            json_match = re.search(r'\{.*\}', plan_text, re.DOTALL)
            if not json_match:
                raise ValueError("在计划响应中未找到JSON")
            plan_data = repair_json(json_match.group(), return_objects=True)

            steps = []
            for step_data in plan_data.get("steps", []):
                tool = step_data.get("tool", "")
                if not tool:
                    tool = ''
                if not isinstance(tool, str):
                    tool = str(tool)

                step = PlanStep(
                    id=step_data.get("id", f"step_{len(steps) + 1}"),
                    description=step_data.get("description", ""),
                    tool=tool,
                    input_template=str(step_data.get("input_template", "")),
                    dependencies=step_data.get("dependencies", []),
                    expected_output=str(step_data.get("expected_output", "")),
                    confidence=step_data.get("confidence", 0.5)
                )
                steps.append(step)

            plan = Plan(
                id=f"plan_{int(time.time())}",
                query=query,
                goal=plan_data.get("goal", ""),
                plan_type=PlanType(plan_data.get("plan_type", "sequential")),
                steps=steps,
                estimated_duration=plan_data.get("estimated_duration", 60.0),
                confidence=plan_data.get("confidence", 0.5),
                metadata={},
                created_at=time.time()
            )
            
            return plan

        except Exception as e:
            # 备选方案：制定一个简单的顺序计划
            logger.error(traceback.format_exc())
            return self._create_fallback_plan(query, available_tools)
    
    def _create_fallback_plan(self, query: str, available_tools: List[Dict]) -> Plan:
        """解析失败时，制定一个简单的顺序计划"""
        import time

        # 简单启发式规则：如果查询涉及计算，使用计算器，如果查询提到搜索/查找，使用搜索工具，以此类推。
        steps = []
        available_tool_names = [tool.get('tool.name') for tool in available_tools]
        if "web_search" in available_tool_names:
            steps.append(PlanStep(
                id="websearch_step",
                description="信息检索",
                tool="web_search",
                input_template=json.dumps({'query': query, 'max_results': 5}, ensure_ascii=False),
                dependencies=[],
                confidence=0.6
            ))

        return Plan(
            id=f"fallback_plan_{int(time.time())}",
            query=query,
            goal="用户提问处理",
            plan_type=PlanType.SEQUENTIAL,
            steps=steps,
            estimated_duration=30.0,
            confidence=0.4,
            metadata={"fallback": True},
            created_at=time.time()
        )
    
    def _apply_refinements(self, plan: Plan, refinement_text: str) -> Plan:
        """改善计划"""
        try:
            json_match = re.search(r'\{.*\}', refinement_text, re.DOTALL)
            if not json_match:
                return plan
            refinements = repair_json(json_match.group(), return_objects=True)

            import copy
            refined_plan = copy.deepcopy(plan)
            refined_plan.id = f"refined_{plan.id}_{int(time.time())}"

            for step_data in refinements.get("add_steps", []):
                tool = step_data.get("tool", "")
                if not tool:
                    tool = ''
                if not isinstance(tool, str):
                    tool = str(tool)

                new_step = PlanStep(
                    id=step_data.get("id"),
                    description=step_data.get("description", ""),
                    tool=tool,
                    input_template=str(step_data.get("input_template", "")),
                    dependencies=step_data.get("dependencies", []),
                    expected_output=str(step_data.get("expected_output")),
                    confidence=step_data.get("confidence", 0.5)
                )
                refined_plan.steps.append(new_step)

            for modification in refinements.get("modify_steps", []):
                step_id = modification.get("id")
                changes = modification.get("changes", {})
                
                for step in refined_plan.steps:
                    if step.id == step_id:
                        for key, value in changes.items():
                            if hasattr(step, key):
                                setattr(step, key, value)

            for step_id in refinements.get("remove_steps", []):
                refined_plan.steps = [s for s in refined_plan.steps if s.id != step_id]

            if "change_plan_type" in refinements:
                refined_plan.plan_type = PlanType(refinements["change_plan_type"])
            
            return refined_plan
            
        except Exception as e:
            return plan
    
    async def _store_plan(self, plan: Plan):
        """保存记忆"""
        self.memory_store.put(
            (self.memory_type, ),
            key = str(uuid.uuid4()),
            value = {
                'query': plan.query,
                'plan': plan.model_dump_json(),
            },
            index=['query']
        )